#include "my-sm4.h"
#include <time.h>

// 字节转化为字
#ifndef GM_GET_UINT32_BE
#define GM_GET_UINT32_BE(n, b, i)                       \
{                                                       \
    (n) = ( (unsigned int) (b)[(i)    ] << 24 )             \
        | ( (unsigned int) (b)[(i) + 1] << 16 )             \
        | ( (unsigned int) (b)[(i) + 2] <<  8 )             \
        | ( (unsigned int) (b)[(i) + 3]       );            \
}
#endif

// 字转化为字节
#ifndef GM_PUT_UINT32_BE
#define GM_PUT_UINT32_BE(n, b ,i)                       \
{                                                       \
    (b)[(i)    ] = (unsigned char) ( (n) >> 24 );       \
    (b)[(i) + 1] = (unsigned char) ( (n) >> 16 );       \
    (b)[(i) + 2] = (unsigned char) ( (n) >>  8 );       \
    (b)[(i) + 3] = (unsigned char) ( (n)       );       \
}
#endif

#define  GM_SHL(x,n) (((x) & 0xFFFFFFFF) << (n % 32))
#define GM_ROTL(x,n) (GM_SHL((x),n) | ((x) >> (32 - (n % 32))))

// S盒 仅用于密钥拓展使用
unsigned char GM_SM4_SBOX[256] = {
	0xd6, 0x90, 0xe9, 0xfe, 0xcc, 0xe1, 0x3d, 0xb7, 0x16, 0xb6, 0x14, 0xc2, 0x28, 0xfb, 0x2c, 0x05,
	0x2b, 0x67, 0x9a, 0x76, 0x2a, 0xbe, 0x04, 0xc3, 0xaa, 0x44, 0x13, 0x26, 0x49, 0x86, 0x06, 0x99,
	0x9c, 0x42, 0x50, 0xf4, 0x91, 0xef, 0x98, 0x7a, 0x33, 0x54, 0x0b, 0x43, 0xed, 0xcf, 0xac, 0x62,
	0xe4, 0xb3, 0x1c, 0xa9, 0xc9, 0x08, 0xe8, 0x95, 0x80, 0xdf, 0x94, 0xfa, 0x75, 0x8f, 0x3f, 0xa6,
	0x47, 0x07, 0xa7, 0xfc, 0xf3, 0x73, 0x17, 0xba, 0x83, 0x59, 0x3c, 0x19, 0xe6, 0x85, 0x4f, 0xa8,
	0x68, 0x6b, 0x81, 0xb2, 0x71, 0x64, 0xda, 0x8b, 0xf8, 0xeb, 0x0f, 0x4b, 0x70, 0x56, 0x9d, 0x35,
	0x1e, 0x24, 0x0e, 0x5e, 0x63, 0x58, 0xd1, 0xa2, 0x25, 0x22, 0x7c, 0x3b, 0x01, 0x21, 0x78, 0x87,
	0xd4, 0x00, 0x46, 0x57, 0x9f, 0xd3, 0x27, 0x52, 0x4c, 0x36, 0x02, 0xe7, 0xa0, 0xc4, 0xc8, 0x9e,
	0xea, 0xbf, 0x8a, 0xd2, 0x40, 0xc7, 0x38, 0xb5, 0xa3, 0xf7, 0xf2, 0xce, 0xf9, 0x61, 0x15, 0xa1,
	0xe0, 0xae, 0x5d, 0xa4, 0x9b, 0x34, 0x1a, 0x55, 0xad, 0x93, 0x32, 0x30, 0xf5, 0x8c, 0xb1, 0xe3,
	0x1d, 0xf6, 0xe2, 0x2e, 0x82, 0x66, 0xca, 0x60, 0xc0, 0x29, 0x23, 0xab, 0x0d, 0x53, 0x4e, 0x6f,
	0xd5, 0xdb, 0x37, 0x45, 0xde, 0xfd, 0x8e, 0x2f, 0x03, 0xff, 0x6a, 0x72, 0x6d, 0x6c, 0x5b, 0x51,
	0x8d, 0x1b, 0xaf, 0x92, 0xbb, 0xdd, 0xbc, 0x7f, 0x11, 0xd9, 0x5c, 0x41, 0x1f, 0x10, 0x5a, 0xd8,
	0x0a, 0xc1, 0x31, 0x88, 0xa5, 0xcd, 0x7b, 0xbd, 0x2d, 0x74, 0xd0, 0x12, 0xb8, 0xe5, 0xb4, 0xb0,
	0x89, 0x69, 0x97, 0x4a, 0x0c, 0x96, 0x77, 0x7e, 0x65, 0xb9, 0xf1, 0x09, 0xc5, 0x6e, 0xc6, 0x84,
	0x18, 0xf0, 0x7d, 0xec, 0x3a, 0xdc, 0x4d, 0x20, 0x79, 0xee, 0x5f, 0x3e, 0xd7, 0xcb, 0x39, 0x48
};

#define GM_SM4_NON_LINEAR_OP(in, out) \
{                                                           \
	(out)  = ( GM_SM4_SBOX[ ((in) >> 24) & 0x0FF ] ) << 24; \
	(out) |= ( GM_SM4_SBOX[ ((in) >> 16) & 0x0FF ] ) << 16; \
	(out) |= ( GM_SM4_SBOX[ ((in) >> 8)  & 0x0FF ] ) << 8;  \
	(out) |= ( GM_SM4_SBOX[ (in)         & 0x0FF ] );       \
}

/* System parameter */
unsigned int GM_SM4_FK[4] = {0xA3B1BAC6, 0x56AA3350, 0x677D9197, 0xB27022DC};

/* fixed parameter */
static const unsigned long GM_SM4_CK[32] =
{
	0x00070e15, 0x1c232a31, 0x383f464d, 0x545b6269,
	0x70777e85, 0x8c939aa1, 0xa8afb6bd, 0xc4cbd2d9,
	0xe0e7eef5, 0xfc030a11, 0x181f262d, 0x343b4249,
	0x50575e65, 0x6c737a81, 0x888f969d, 0xa4abb2b9,
	0xc0c7ced5, 0xdce3eaf1, 0xf8ff060d, 0x141b2229,
	0x30373e45, 0x4c535a61, 0x686f767d, 0x848b9299,
	0xa0a7aeb5, 0xbcc3cad1, 0xd8dfe6ed, 0xf4fb0209,
	0x10171e25, 0x2c333a41, 0x484f565d, 0x646b7279
};

static void sm4_key_schedule(const unsigned char * key, unsigned int * rk)
{
	unsigned int tmp1, tmp2, K[36];
	int i;

	// transfer input
	// K0, K1, K2, K3
    GM_GET_UINT32_BE( K[0], key, 0);
    GM_GET_UINT32_BE( K[1], key, 4);
    GM_GET_UINT32_BE( K[2], key, 8);
    GM_GET_UINT32_BE( K[3], key, 12);

    K[0] ^= GM_SM4_FK[0];
    K[1] ^= GM_SM4_FK[1];
    K[2] ^= GM_SM4_FK[2];
    K[3] ^= GM_SM4_FK[3];

	// rki = Ki+4
	for(i = 0; i < 32; i++)
	{
		// non linear's input
		tmp1 = K[i + 1] ^ K[i + 2] ^ K[i + 3] ^ GM_SM4_CK[i];

		// non linear operation
		GM_SM4_NON_LINEAR_OP(tmp1, tmp2);

		// linear operation
		rk[i] = K[i + 4] = K[i] ^ (tmp2 ^ GM_ROTL(tmp2, 13) ^ GM_ROTL(tmp2, 23));
	}
}

static void KeyRevser(unsigned int * rk, unsigned int * dec_rk) {
    for (int i = 32 - 1, j = 0; i >= 0 ; i --, j++) {
        dec_rk[j] = rk[i];
    }
}

typedef struct {
    unsigned int X1;
    unsigned int X2;
    unsigned int X3;
    unsigned int X4;
}BLOCK;

void Reverse(BLOCK * target) {
    unsigned int tmp;
    for(int i = 0; i < 256/sizeof(BLOCK); i++ ) {
        tmp = target[i].X1;
        target[i].X1 = target[i].X4;
        target[i].X4 = tmp;

        tmp = target[i].X2;
        target[i].X2 = target[i].X3;
        target[i].X3 = tmp;
    }
}


static void F256(unsigned char *plaintext, unsigned int * rk) {
    // 每次处理256字节的数据
    BLOCK * block = (BLOCK *)plaintext;
    unsigned int buf[64] = {0};
    
    // 32轮函数
    for(int r = 0; r < 32; r ++) {
        for(int i = 0; i < 256 / sizeof(BLOCK); i++) {
            buf[i] = block[i].X2 ^ block[i].X3 ^ block[i].X4 ^ rk[r];
        }


        //search sbox
        long_bits in = BUILDLONGBITS();
        long_bits out = BUILDLONGBITS();

        Block2LongBits((unsigned char *)buf, &in);
        LongBitsSbox(in, &out);
        LongBits2Block(out, (unsigned char *)buf);

        //ROL and xor
        for(int i = 0; i < 256 / sizeof(BLOCK); i++){
            buf[i] = block[i].X1 ^ buf[i] ^ (GM_ROTL(buf[i], 2)) ^ (GM_ROTL(buf[i], 10)) ^ (GM_ROTL(buf[i], 18)) ^ (GM_ROTL(buf[i], 24));
        }

        // change the plaintext
        memcpy(plaintext, plaintext + sizeof(unsigned int), 256 - sizeof(unsigned int));

        for(int i = 0; i < 256 / sizeof(BLOCK); i++) {
            block[i].X4 = buf[i];
        }
        memset(buf, 0, 256);
    }

    // revser
    Reverse(block);

}

void BlockEndianTrans(unsigned char * target) {
    unsigned char tmp;
    for(int i = 0; i < 256; i+=4) {
        tmp = target[i];
        target[i] = target[i + 3];
        target[i+3] = tmp;

        tmp = target[i + 1];
        target[i + 1] = target[i + 2];
        target[i+2] = tmp;

    }
}

unsigned char * EncryptECB(unsigned char * key, unsigned char * plaintext, int len) {

    unsigned int rk[32];
    sm4_key_schedule(key, rk);

    for(int i = 0; i < len / 256; i++) {
        F256(plaintext + 256*i, rk);
        
        //不知为何，这个算法下得到的加密数据大小端与实际的sm4输出是相反的，需要一个大小端的转换操作
        BlockEndianTrans(plaintext + i*256);
    }

    return plaintext;
}

unsigned char * DecryptECB(unsigned char * key, unsigned char * cipher, int len) {
    unsigned int rk[32];
    sm4_key_schedule(key, rk);

    unsigned int dec_rk[32];
    KeyRevser(rk, dec_rk);

    for(int i = 0; i < len / 256; i++) {
        BlockEndianTrans(cipher + i*256);
        F256(cipher + 256*i, dec_rk); 
    }

    return cipher;

}


typedef struct {
    unsigned char key[16];
    unsigned int rk[32];
    unsigned int dec_key[32];
    unsigned char * (*ENCRYPT)(unsigned char * , unsigned char * , int);
    unsigned char * (*DECRYPT)(unsigned char * , unsigned char * , int);
}SM4_ECB;

void BUILDSM4_ECB(SM4_ECB * sm4_ctx, unsigned char * key) { 
    memcpy(sm4_ctx->rk, key, 16);
    sm4_key_schedule(sm4_ctx->key, sm4_ctx->rk);
    KeyRevser(sm4_ctx->rk, sm4_ctx->dec_key);
    sm4_ctx->ENCRYPT = &EncryptECB;
    sm4_ctx->DECRYPT = &DecryptECB;
}
#define PLAINTEXTBLOCKSIZE 256

int test() {
    //128:32k   *64 : 16k    32:8K
    unsigned char plaintext[PLAINTEXTBLOCKSIZE];
    unsigned char * tmp = "SM4 is a symmetric block cipher algorithm. It employs a 32-round Feistel network structure, combining operations such as substitution, and permutation, and XOR. SM4 is used for data protection and confidentiality in various applications due to robustness.";
    //memcpy(plaintext, tmp, 256);
    for(unsigned int i=0;i<1;i++){
        memcpy(plaintext+i*PLAINTEXTBLOCKSIZE, tmp, PLAINTEXTBLOCKSIZE);
    }
    printf("\n================================================================================================================================================================= \n");
    clock_t begin, end;
    printf("\nmessage: \n  ");
    for(int i = 0; i < 256; i++) {
        printf("%c", plaintext[i]);
    }
    printf("\n");
    
    unsigned int rk[32];
    unsigned char key[] = {0x11, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x10};
    sm4_key_schedule(key, rk);
    printf("\nkey: \n  ");
    for(int i = 0; i < 32; i++) {
        printf("%2x ", key[i]);
    }
    printf("\n");
    begin  = clock();
    EncryptECB(key, plaintext, PLAINTEXTBLOCKSIZE);
    end = clock();
    double d1 = (double)(end-begin)*1000/CLOCKS_PER_SEC;
    printf("\nciphertext: \n  ");
    for(int i = 0; i < 256; i++) {
        printf("%2x", plaintext[i]);
        if((i+1) % 64 == 0){
            printf("\n  ");
        }
    }
    // puts("");
    begin  = clock();
    DecryptECB(key, plaintext, PLAINTEXTBLOCKSIZE);
    end = clock();
    double d2 = (double)(end-begin)*1000/CLOCKS_PER_SEC;

    printf("plaintext: \n  ");
    for(int i = 0; i < 256; i++) {
        printf("%c", plaintext[i]);
    }
    // puts("");
    printf("\n \n");
    printf("encryption time: %f ms \n", d1);
    printf("decryption time: %f ms \n", d2); 
    printf("\n================================================================================================================================================================= \n");


    return 0;
}

int main() {
    test();
    return 0;
}


